import os
import sys
import requests
from urllib.parse import urljoin
from concurrent.futures import ThreadPoolExecutor
from Crypto.Cipher import AES
import shutil
import re
import binascii
import logging
import subprocess
import time
import argparse

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def download_m3u8(m3u8_url):
    try:
        response = requests.get(m3u8_url)
        response.raise_for_status()
        m3u8_content = response.text
        with open('playlist.m3u8', 'w') as f:
            f.write(m3u8_content)
        logging.info("已下载m3u8文件")
        return m3u8_content
    except requests.exceptions.RequestException as e:
        logging.error(f"下载m3u8文件失败: {e}")
        sys.exit(1)

def parse_m3u8(m3u8_content, base_url):
    ts_urls = []
    key = None
    iv = None
    lines = m3u8_content.strip().split('\n')
    for index, line in enumerate(lines):
        if line.startswith("#EXT-X-KEY"):
            method = re.search(r'METHOD=([A-Za-z0-9-]+)', line)
            uri = re.search(r'URI="([^"]+)"', line)
            iv_match = re.search(r'IV=0x([0-9a-fA-F]+)', line)

            if method and method.group(1) != 'AES-128':
                raise ValueError(f"不支持的加密方法: {method.group(1)}")
            if uri:
                key_url = uri.group(1)
                key_url = urljoin(base_url, key_url)
                key = download_key(key_url)
                logging.info("已下载并保存key")
            if iv_match:
                iv_hex = iv_match.group(1)
                iv = binascii.unhexlify(iv_hex)
                logging.info(f"已解析IV: {iv_hex}")
            else:
                iv = None  # 使用默认IV
        elif line and not line.startswith("#"):
            ts_url = urljoin(base_url, line)
            ts_urls.append(ts_url)
    return ts_urls, key, iv

def download_key(key_url):
    try:
        response = requests.get(key_url)
        response.raise_for_status()
        with open('key.key', 'wb') as f:
            f.write(response.content)
        return response.content
    except requests.exceptions.RequestException as e:
        logging.error(f"下载key失败: {e}")
        sys.exit(1)

def download_ts(url, index, key, iv, ts_folder):
    try:
        ts_content = requests.get(url).content
        if key:
            if iv is None:
                iv = index.to_bytes(16, byteorder='big')  # 默认IV
            cipher = AES.new(key, AES.MODE_CBC, iv=iv)
            ts_content = cipher.decrypt(ts_content)
        ts_filename = os.path.join(ts_folder, f"{index}.ts")
        os.makedirs(ts_folder, exist_ok=True)  # 确保文件夹存在
        with open(ts_filename, 'wb') as f:
            f.write(ts_content)
        logging.info(f"已下载并保存{ts_filename}")
    except requests.exceptions.RequestException as e:
        logging.error(f"下载TS文件失败: {e}")
    except Exception as e:
        logging.error(f"解密TS文件失败: {e}")

def merge_ts_files(num_files, ts_folder, output_filename):
    try:
        with open(output_filename, 'wb') as merged:
            logging.info(f"开始合并所有.ts文件为{output_filename}")
            for i in range(num_files):
                ts_filename = os.path.join(ts_folder, f"{i}.ts")
                if os.path.exists(ts_filename):
                    with open(ts_filename, 'rb') as mergefile:
                        shutil.copyfileobj(mergefile, merged)
                    logging.info(f"已合并 {ts_filename}")
                else:
                    logging.warning(f"文件 {ts_filename} 不存在，跳过")
        logging.info(f"所有可用的.ts文件已合并为{output_filename}")
    except Exception as e:
        logging.error(f"合并TS文件失败: {e}")

def convert_to_mp4(output_filename):
    try:
        mp4_filename = output_filename.replace(".ts", ".mp4")
        logging.info(f"开始转换为MP4格式：{mp4_filename}")
        subprocess.run(['ffmpeg', '-i', output_filename, '-c:v', 'libx264', '-c:a', 'aac', mp4_filename], check=True)
        logging.info(f"已转换为MP4格式: {mp4_filename}")
        # 转换成功后，删除ts文件
        os.remove(output_filename)  
        logging.info(f"已删除合并后的TS文件: {output_filename}")
    except subprocess.CalledProcessError as e:
        logging.error(f"转换为MP4失败: {e}")

def cleanup_ts_files(ts_folder):
    try:
        for filename in os.listdir(ts_folder):
            file_path = os.path.join(ts_folder, filename)
            if os.path.isfile(file_path):
                os.remove(file_path)
        os.rmdir(ts_folder)
        logging.info("已清理临时TS文件")
    except Exception as e:
        logging.error(f"清理文件夹失败: {e}")

def main(m3u8_url, ts_folder, output_filename):
    base_url = os.path.dirname(m3u8_url) + '/'
    m3u8_content = download_m3u8(m3u8_url)
    ts_urls, key, iv = parse_m3u8(m3u8_content, base_url)

    with ThreadPoolExecutor(max_workers=5) as executor:
        for index, ts_url in enumerate(ts_urls):
            executor.submit(download_ts, ts_url, index, key, iv, ts_folder)

    merge_ts_files(len(ts_urls), ts_folder, output_filename)
    convert_to_mp4(output_filename)  # 转换为MP4
    cleanup_ts_files(ts_folder)  # 清理临时的TS文件夹

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='下载并合并m3u8视频流，最终输出MP4文件')
    parser.add_argument('m3u8_url', help='m3u8文件的URL')
    parser.add_argument('--ts_folder', default=time.strftime("%Y-%m-%d"), help='TS文件存储文件夹（默认为当天日期）')
    parser.add_argument('--output_filename', default=time.strftime("%Y-%m-%d") + ".ts", help='合并后的TS输出文件名（默认为当天日期）')
    
    args = parser.parse_args()

    main(args.m3u8_url, args.ts_folder, args.output_filename)